import torch
import torch.nn as nn
import math
from torch_geometric.nn import global_mean_pool


# RNN model for vulnerability detection
# class RNN(nn.Module):
#     def __init__(self, nfeat, nhid, nclass, nlayers=6, device=None, dropout=0.5):
#         super(RNN, self).__init__()
#         self.device, self.nfeat, self.nhid, self.nclass, self.nlayers = device, nfeat, nhid, nclass, nlayers
#         self.rnn = nn.RNN(nfeat, nhid, nlayers, )
#         self.linear = nn.Linear(nhid, nclass)
#         self.sigmoid = nn.Sigmoid()
#         if device is None:
#             self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#         else:
#             self.device = device


#     def forward(self, data):
#         data = data.to(self.device)
#         # 初始化隐藏状态
#         h0 = torch.zeros(self.num_layers, data.size(0), self.hidden_dim)   
#         # RNN前向传播
#         out, _ = self.rnn(data, h0)
#         # 取最后一个时间步的输出
#         out = out[:, -1, :]
        
#         # 全连接层和激活函数
#         out = self.linear(out)
#         out = self.sigmoid(out)
#         return out

    
# if __name__ == '__main__':

#     pass